package com.mcp.client;

import com.mcp.client.transport.*;
import com.mcp.server.model.McpRequest;
import com.mcp.server.model.McpResponse;
import com.mcp.util.JsonUtils;

import java.io.IOException;
import java.net.ServerSocket;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * 自动选择最佳通信方式的MCP客户端
 */
public class AutoSelectMcpClient {
    private static final String MCP_STDIO_ENV = "MCP_STDIO";
    private static final int HTTP_SSE_PORT = 8081;
    
    private McpTransport transport;
    private boolean useAsync = false;
    
    public AutoSelectMcpClient() throws Exception {
        this.transport = selectBestTransport();
        this.useAsync = (transport instanceof HttpSseTransport);
    }
    
    /**
     * 选择最佳传输方式
     */
    private McpTransport selectBestTransport() throws Exception {
        System.out.println("[MCP Client] 开始选择最佳通信方式...");
        
        // 1. 检查环境变量 → 存在MCP_STDIO？ → Stdio（最高优先级）
        if (checkStdioEnvironment()) {
            System.out.println("[MCP Client] 使用 Stdio 通信方式 (最高优先级)");
            return new StdioTransport();
        }
        
        // 2. 检测端口可用性 → 能绑定本地端口？ → HTTP+SSE
        if (checkHttpSsePortAvailable()) {
            System.out.println("[MCP Client] 使用 HTTP+SSE 通信方式");
            return new HttpSseTransport();
        }
        
        // 3. 默认回退 → Streamable HTTP（最通用）
        System.out.println("[MCP Client] 使用 Streamable HTTP 通信方式 (默认回退)");
        return new StreamableHttpTransport();
    }
    
    /**
     * 检查是否存在MCP_STDIO环境变量
     */
    private boolean checkStdioEnvironment() {
        String envValue = System.getenv(MCP_STDIO_ENV);
        return envValue != null && !envValue.trim().isEmpty();
    }
    
    /**
     * 检查HTTP+SSE端口是否可用
     */
    private boolean checkHttpSsePortAvailable() {
        try (ServerSocket socket = new ServerSocket(HTTP_SSE_PORT)) {
            // 端口可用，但我们不需要保持这个连接
            return true;
        } catch (IOException e) {
            // 端口已被占用，可能有服务正在运行
            return false;
        }
    }
    
    /**
     * 调用远程方法
     */
    public McpResponse call(String method, Object params) throws Exception {
        McpRequest request = new McpRequest();
        request.setMethod(method);
        request.setParams(params);
        request.setId(System.currentTimeMillis()); // 使用时间戳作为ID

        String requestJson = JsonUtils.toJson(request);
        System.out.println("客户端发送请求: " + requestJson);

        String responseJson = transport.send(requestJson);
        System.out.println("客户端接收到响应: " + responseJson);

        return JsonUtils.fromJson(responseJson, McpResponse.class);
    }
    
    /**
     * 异步调用远程方法（仅当使用HTTP+SSE时可用）
     */
    public CompletableFuture<McpResponse> callAsync(String method, Object params) throws Exception {
        if (!(transport instanceof HttpSseTransport)) {
            // 如果不是HTTP+SSE传输，使用同步调用并包装为异步
            return CompletableFuture.supplyAsync(() -> {
                try {
                    return call(method, params);
                } catch (Exception e) {
                    throw new RuntimeException("异步调用失败", e);
                }
            });
        }
        
        HttpSseTransport httpTransport = (HttpSseTransport) transport;
        CompletableFuture<McpResponse> future = new CompletableFuture<>();
        
        // 生成唯一请求ID
        String requestId = UUID.randomUUID().toString();
        
        // 注册响应处理器
        httpTransport.registerResponseHandler(requestId, response -> {
            future.complete(response);
        });
        
        // 发送请求
        McpRequest request = new McpRequest();
        request.setMethod(method);
        request.setParams(params);
        request.setId(requestId);
        
        String requestJson = JsonUtils.toJson(request);
        System.out.println("客户端异步发送请求: " + requestJson);
        
        // 启动请求线程
        new Thread(() -> {
            try {
                transport.send(requestJson);
            } catch (Exception e) {
                future.completeExceptionally(e);
            }
        }).start();
        
        return future;
    }
    
    /**
     * 使用流式处理响应
     */
    public void callWithStreaming(String method, Object params, Consumer<McpResponse> responseHandler) throws Exception {
        if (transport instanceof StreamableHttpTransport) {
            McpRequest request = new McpRequest();
            request.setMethod(method);
            request.setParams(params);
            request.setId(System.currentTimeMillis());
            
            String requestJson = JsonUtils.toJson(request);
            System.out.println("客户端发送流式请求: " + requestJson);
            
            ((StreamableHttpTransport) transport).sendWithStreaming(requestJson, responseHandler);
        } else {
            // 如果不是StreamableHttpTransport，使用普通调用
            McpResponse response = call(method, params);
            responseHandler.accept(response);
        }
    }
    
    /**
     * 关闭客户端
     */
    public void close() throws Exception {
        if (transport != null) {
            transport.close();
        }
    }
    
    /**
     * 示例用法
     */
    public static void main(String[] args) {
        try {
            // 创建自动选择最佳通信方式的客户端
            AutoSelectMcpClient client = new AutoSelectMcpClient();
            
            // 调用add方法
            Map<String, Integer> mathParams = new HashMap<>();
            mathParams.put("a", 5);
            mathParams.put("b", 3);
            
            // 同步调用
            McpResponse response = client.call("add", mathParams);
            System.out.println("5 + 3 = " + response.getResult());
            
            // 异步调用（如果支持）
            // if (client instanceof HttpSseTransport) {
                CompletableFuture<McpResponse> asyncResponse = client.callAsync("add", mathParams);
                System.out.println("异步调用已发送，等待响应...");
                
                // 等待异步响应
                asyncResponse.thenAccept(result -> {
                    System.out.println("异步调用结果: 5 + 3 = " + result.getResult());
                }).exceptionally(e -> {
                    System.err.println("异步调用失败: " + e.getMessage());
                    return null;
                });
                
                // 等待异步操作完成
                TimeUnit.SECONDS.sleep(2);
            // }
            
            // 流式调用（如果支持）
            client.callWithStreaming("add", mathParams, streamResponse -> {
                System.out.println("流式响应: 5 + 3 = " + streamResponse.getResult());
            });
            
            // 关闭客户端
            client.close();
            
        } catch (Exception e) {
            System.err.println("调用失败: " + e.getMessage());
            e.printStackTrace();
        }
    }
}